{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import modules that would be used later.\n",
    "import os, random, pickle, itertools, copy\n",
    "import learn2learn as l2l\n",
    "import torchvision as tv\n",
    "from PIL.Image import LANCZOS\n",
    "from learn2learn.data.transforms import NWays, KShots, LoadData, RemapLabels, ConsecutiveLabels\n",
    "from learn2learn.vision.transforms import RandomClassRotation\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([5, 1, 28, 28])\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# 1. Apply transforms on input data if any (eg., Resizing, converting arrays to tensors)\n",
    "data_transform = tv.transforms.Compose([tv.transforms.Resize((28, 28), interpolation=LANCZOS), tv.transforms.ToTensor()]) \n",
    "\n",
    "# 2. Download the dataset, and apply the above transforms on each of the samples\n",
    "dataset = tv.datasets.MNIST('/root/datasets/', train=False, transform=data_transform)\n",
    "\n",
    "# 3. Wrap the dataset using MetaDataset to enable fast indexing\n",
    "omniglot = l2l.data.MetaDataset(dataset) # Use MetaDataset to do fast indexing of samples\n",
    "\n",
    "# 4. Specify transforms to be used for generating tasks\n",
    "transforms = [\n",
    "                    NWays(omniglot, 5),  # Samples N random classes per task (here, N = 5)\n",
    "                    KShots(omniglot, 1), # Samples K samples per class from the above N classes (here, K = 1) \n",
    "                    LoadData(omniglot), # Loads a sample from the dataset\n",
    "                    RemapLabels(omniglot), # Remaps labels starting from zero\n",
    "                    ConsecutiveLabels(omniglot), # Re-orders samples s.t. they are sorted in consecutive order \n",
    "                    RandomClassRotation(omniglot, [0, 90, 180, 270]) # Randomly rotate sample over x degrees (only for vision tasks)\n",
    "                    ]\n",
    "\n",
    "# 5. Generate set of tasks using the dataset, and transforms\n",
    "taskset = l2l.data.TaskDataset(dataset=omniglot, task_transforms=transforms, num_tasks=10) # Creates sets of tasks from the dataset \n",
    "\n",
    "# Now sample a task from the taskset\n",
    "X, y = taskset.sample()\n",
    "# or, you can also sample this way:\n",
    "for task in taskset:\n",
    "    X, y = task\n",
    "print(X.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([3, 0, 1, 2, 4])\n",
      "tensor([3, 2, 0, 1, 4])\n",
      "tensor([2, 0, 4, 1, 3])\n",
      "tensor([0, 4, 2, 1, 3])\n",
      "tensor([2, 4, 1, 0, 3])\n",
      "tensor([1, 0, 2, 3, 4])\n",
      "tensor([4, 3, 1, 2, 0])\n",
      "tensor([1, 3, 0, 2, 4])\n",
      "tensor([4, 0, 2, 1, 3])\n",
      "tensor([0, 1, 4, 2, 3])\n"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "for task in taskset:\n",
    "    X, y = task\n",
    "    print(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl4",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
